# -*- coding: utf-8 -*-
"""
简化的对比分析和可视化
分别分析embedding+rerank和rerank only两种实验的结果
"""

import os
import json
import time
from typing import List, Dict, Any, Optional
from dataclasses import dataclass
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from datetime import datetime
from loguru import logger

from .embedding_rerank_comparison import EmbeddingRerankComparison, EmbeddingRerankResult
from .rerank_only_comparison import RerankOnlyComparison, RerankOnlyResult
from ...config import (
    get_rerank_experiment_config, RERANK_EVALUATION_CONFIG,
    RERANK_VISUALIZATION_CONFIG, DEFAULT_DEVICE,
    MODELS_CACHE_DIR, RESULTS_CACHE_DIR
)


class SimplifiedComparison:
    """
    简化的对比分析类
    
    分别分析embedding+rerank和rerank only两种实验的结果
    """
    
    def __init__(self, output_dir: str = "reports/rerank"):
        """
        初始化综合对比分析
        
        Args:
            output_dir: 输出目录
        """
        self.output_dir = output_dir
        self.embedding_rerank_results = []
        self.rerank_only_results = []
        
        # 创建输出目录
        os.makedirs(output_dir, exist_ok=True)
        
        logger.info(f"初始化简化对比分析，输出目录: {output_dir}")
    
    def add_embedding_rerank_results(self, results: List[EmbeddingRerankResult]) -> None:
        """
        添加embedding+rerank实验结果
        
        Args:
            results: embedding+rerank实验结果列表
        """
        self.embedding_rerank_results.extend(results)
        logger.info(f"添加了 {len(results)} 个embedding+rerank实验结果")
    
    def add_rerank_only_results(self, results: List[RerankOnlyResult]) -> None:
        """
        添加rerank only实验结果
        
        Args:
            results: rerank only实验结果列表
        """
        self.rerank_only_results.extend(results)
        logger.info(f"添加了 {len(results)} 个rerank only实验结果")
    
    def generate_analysis_report(self) -> Dict[str, Any]:
        """
        生成分析报告
        
        Returns:
            Dict[str, Any]: 分析报告
        """
        report = {
            'metadata': {
                'generated_at': datetime.now().isoformat(),
                'embedding_rerank_count': len(self.embedding_rerank_results),
                'rerank_only_count': len(self.rerank_only_results)
            },
            'embedding_rerank_analysis': {},
            'rerank_only_analysis': {}
        }
        
        # 分析embedding+rerank结果
        if self.embedding_rerank_results:
            report['embedding_rerank_analysis'] = self._analyze_embedding_rerank_results()
        
        # 分析rerank only结果
        if self.rerank_only_results:
            report['rerank_only_analysis'] = self._analyze_rerank_only_results()
        
        return report
    
    def _analyze_embedding_rerank_results(self) -> Dict[str, Any]:
        """
        分析embedding+rerank结果
        """
        # 按模型分组
        model_results = {}
        for result in self.embedding_rerank_results:
            model_name = result.model_name
            if model_name not in model_results:
                model_results[model_name] = []
            model_results[model_name].append(result)
        
        analysis = {}
        for model_name, results in model_results.items():
            total_times = [r.total_time for r in results]
            rerank_times = [r.rerank_time for r in results]
            embedding_times = [r.embedding_time for r in results]
            
            # 收集rerank分数
            all_rerank_scores = []
            for r in results:
                all_rerank_scores.extend(r.rerank_scores)
            
            analysis[model_name] = {
                'performance': {
                    'avg_total_time': np.mean(total_times),
                    'std_total_time': np.std(total_times),
                    'avg_rerank_time': np.mean(rerank_times),
                    'avg_embedding_time': np.mean(embedding_times),
                    'rerank_time_ratio': np.mean(rerank_times) / np.mean(total_times)
                },
                'quality': {
                    'avg_rerank_score': np.mean(all_rerank_scores) if all_rerank_scores else 0,
                    'std_rerank_score': np.std(all_rerank_scores) if all_rerank_scores else 0,
                    'score_range': {
                        'min': np.min(all_rerank_scores) if all_rerank_scores else 0,
                        'max': np.max(all_rerank_scores) if all_rerank_scores else 0
                    }
                },
                'sample_count': len(results)
            }
        
        return analysis
    
    def _analyze_rerank_only_results(self) -> Dict[str, Any]:
        """
        分析rerank only结果
        """
        # 按模型分组
        model_results = {}
        for result in self.rerank_only_results:
            model_name = result.model_name
            if model_name not in model_results:
                model_results[model_name] = []
            model_results[model_name].append(result)
        
        analysis = {}
        for model_name, results in model_results.items():
            processing_times = [r.processing_time for r in results]
            score_variances = [r.score_variance for r in results]
            
            # 收集top分数
            all_top_scores = []
            for r in results:
                all_top_scores.extend(r.top_scores)
            
            analysis[model_name] = {
                'performance': {
                    'avg_processing_time': np.mean(processing_times),
                    'std_processing_time': np.std(processing_times),
                    'min_processing_time': np.min(processing_times),
                    'max_processing_time': np.max(processing_times)
                },
                'quality': {
                    'avg_score_variance': np.mean(score_variances),
                    'avg_top_score': np.mean(all_top_scores) if all_top_scores else 0,
                    'std_top_score': np.std(all_top_scores) if all_top_scores else 0,
                    'score_consistency': 1.0 / (1.0 + np.mean(score_variances))  # 一致性指标
                },
                'sample_count': len(results)
            }
        
        return analysis
    

    

    

    
    def save_analysis_report(self) -> str:
        """
        保存分析报告
        
        Returns:
            str: 报告文件路径
        """
        report = self.generate_analysis_report()
        
        report_path = os.path.join(self.output_dir, 'analysis_report.json')
        with open(report_path, 'w', encoding='utf-8') as f:
            json.dump(report, f, ensure_ascii=False, indent=2)
        
        logger.info(f"分析报告已保存到: {report_path}")
        return report_path
    
    def create_visualization(self) -> None:
        """
        创建可视化图表
        """
        if not self.embedding_rerank_results and not self.rerank_only_results:
            logger.warning("没有结果可以可视化")
            return
        
        # SCI paper style settings
        plt.rcParams.update({
            'font.family': 'serif',
            'font.serif': ['Times New Roman', 'DejaVu Serif'],
            'font.size': 12,
            'axes.linewidth': 1.5,
            'axes.spines.top': False,
            'axes.spines.right': False,
            'axes.grid': True,
            'grid.alpha': 0.3,
            'grid.linewidth': 0.8,
            'legend.frameon': False,
            'legend.fontsize': 11,
            'xtick.direction': 'in',
            'ytick.direction': 'in',
            'xtick.major.size': 5,
            'ytick.major.size': 5,
            'xtick.minor.size': 3,
            'ytick.minor.size': 3,
            'axes.labelsize': 12,
            'axes.titlesize': 14,
            'figure.titlesize': 16
        })
        
        # 创建分别的可视化
        if self.embedding_rerank_results:
            self._create_embedding_rerank_visualization()
        
        if self.rerank_only_results:
            self._create_rerank_only_visualization()
    
    def _create_embedding_rerank_visualization(self) -> None:
        """
        创建embedding+rerank可视化
        """
        fig, axes = plt.subplots(2, 2, figsize=(14, 10))
        fig.suptitle('Embedding+Rerank Model Analysis', fontweight='bold')
        
        # 1. 时间分布
        self._plot_embedding_rerank_time_distribution(axes[0, 0])
        
        # 2. 分数分布
        self._plot_embedding_rerank_score_distribution(axes[0, 1])
        
        # 3. 模型性能对比
        self._plot_embedding_rerank_performance(axes[1, 0])
        
        # 4. 时间组成分析
        self._plot_embedding_rerank_time_breakdown(axes[1, 1])
        
        plt.tight_layout()
        
        output_path = os.path.join(self.output_dir, 'embedding_rerank_analysis.png')
        plt.savefig(output_path, dpi=300, bbox_inches='tight', 
                   facecolor='white', edgecolor='none')
        plt.close()
        
        logger.info(f"Embedding+Rerank可视化已保存到: {output_path}")
    
    def _create_rerank_only_visualization(self) -> None:
        """
        创建rerank only可视化
        """
        fig, axes = plt.subplots(2, 2, figsize=(14, 10))
        fig.suptitle('Rerank Only Model Analysis', fontweight='bold')
        
        # 1. 时间分布
        self._plot_rerank_only_time_distribution(axes[0, 0])
        
        # 2. 分数分布
        self._plot_rerank_only_score_distribution(axes[0, 1])
        
        # 3. 模型性能对比
        self._plot_rerank_only_performance(axes[1, 0])
        
        # 4. 分数一致性分析
        self._plot_rerank_only_consistency(axes[1, 1])
        
        plt.tight_layout()
        
        output_path = os.path.join(self.output_dir, 'rerank_only_analysis.png')
        plt.savefig(output_path, dpi=300, bbox_inches='tight', 
                   facecolor='white', edgecolor='none')
        plt.close()
        
        logger.info(f"Rerank Only可视化已保存到: {output_path}")
    
    def _plot_embedding_rerank_time_distribution(self, ax):
        """绘制embedding+rerank时间分布"""
        sci_colors = ['#2E86AB', '#A23B72', '#F18F01', '#C73E1D']
        
        # 按模型分组
        model_data = {}
        for result in self.embedding_rerank_results:
            model_name = result.model_name
            if model_name not in model_data:
                model_data[model_name] = []
            model_data[model_name].append(result.total_time)
        
        # 绘制箱线图
        if model_data:
            models = list(model_data.keys())
            data = [model_data[model] for model in models]
            
            bp = ax.boxplot(data, labels=models, patch_artist=True)
            
            for patch, color in zip(bp['boxes'], sci_colors[:len(bp['boxes'])]):
                patch.set_facecolor(color)
                patch.set_alpha(0.7)
                patch.set_edgecolor('black')
                patch.set_linewidth(1.0)
            
            ax.set_title('Total Time Distribution', fontweight='bold')
            ax.set_ylabel('Time (seconds)')
            ax.tick_params(axis='x', rotation=45)
            ax.grid(True, alpha=0.3)
    
    def _plot_embedding_rerank_score_distribution(self, ax):
        """绘制embedding+rerank分数分布"""
        sci_colors = ['#2E86AB', '#A23B72', '#F18F01', '#C73E1D']
        
        # 按模型收集分数
        model_scores = {}
        for result in self.embedding_rerank_results:
            model_name = result.model_name
            if model_name not in model_scores:
                model_scores[model_name] = []
            model_scores[model_name].extend(result.rerank_scores)
        
        # 绘制分数分布
        for i, (model, scores) in enumerate(model_scores.items()):
            if scores:
                ax.hist(scores, alpha=0.6, label=model, bins=20, density=True,
                       color=sci_colors[i % len(sci_colors)], edgecolor='black', linewidth=0.5)
        
        ax.set_title('Rerank Score Distribution', fontweight='bold')
        ax.set_xlabel('Score')
        ax.set_ylabel('Density')
        ax.legend()
        ax.grid(True, alpha=0.3)
    
    def _plot_embedding_rerank_performance(self, ax):
        """绘制embedding+rerank性能对比"""
        sci_colors = ['#2E86AB', '#A23B72', '#F18F01', '#C73E1D']
        
        # 按模型统计性能
        model_stats = {}
        for result in self.embedding_rerank_results:
            model_name = result.model_name
            if model_name not in model_stats:
                model_stats[model_name] = {'times': [], 'scores': []}
            model_stats[model_name]['times'].append(result.total_time)
            model_stats[model_name]['scores'].extend(result.rerank_scores)
        
        models = list(model_stats.keys())
        avg_times = [np.mean(model_stats[model]['times']) for model in models]
        avg_scores = [np.mean(model_stats[model]['scores']) if model_stats[model]['scores'] else 0 for model in models]
        
        bars = ax.bar(models, avg_scores, color=sci_colors[:len(models)], alpha=0.7, edgecolor='black')
        
        # 添加时间信息作为标签
        for bar, time in zip(bars, avg_times):
            height = bar.get_height()
            ax.text(bar.get_x() + bar.get_width()/2., height + 0.01,
                   f'{time:.3f}s', ha='center', va='bottom', fontsize=9)
        
        ax.set_title('Model Performance Comparison', fontweight='bold')
        ax.set_ylabel('Average Rerank Score')
        ax.tick_params(axis='x', rotation=45)
        ax.grid(True, alpha=0.3)
    
    def _plot_embedding_rerank_time_breakdown(self, ax):
        """绘制embedding+rerank时间组成分析"""
        sci_colors = ['#2E86AB', '#A23B72']
        
        # 按模型统计时间组成
        model_times = {}
        for result in self.embedding_rerank_results:
            model_name = result.model_name
            if model_name not in model_times:
                model_times[model_name] = {'embedding': [], 'rerank': []}
            model_times[model_name]['embedding'].append(result.embedding_time)
            model_times[model_name]['rerank'].append(result.rerank_time)
        
        models = list(model_times.keys())
        embedding_times = [np.mean(model_times[model]['embedding']) for model in models]
        rerank_times = [np.mean(model_times[model]['rerank']) for model in models]
        
        x = np.arange(len(models))
        width = 0.35
        
        ax.bar(x - width/2, embedding_times, width, label='Embedding Time', 
               color=sci_colors[0], alpha=0.7, edgecolor='black')
        ax.bar(x + width/2, rerank_times, width, label='Rerank Time', 
               color=sci_colors[1], alpha=0.7, edgecolor='black')
        
        ax.set_title('Time Breakdown Analysis', fontweight='bold')
        ax.set_ylabel('Time (seconds)')
        ax.set_xticks(x)
        ax.set_xticklabels(models, rotation=45)
        ax.legend()
        ax.grid(True, alpha=0.3)
    
    def _plot_rerank_only_time_distribution(self, ax):
        """绘制rerank only时间分布图"""
        # SCI color scheme
        sci_colors = ['#2E86AB', '#A23B72', '#F18F01', '#C73E1D']
        
        if not self.rerank_only_results:
            ax.text(0.5, 0.5, 'No Data Available', ha='center', va='center', 
                   transform=ax.transAxes, fontsize=11)
            return
        
        # 按模型分组收集时间数据
        model_times = {}
        for result in self.rerank_only_results:
            model_name = result.model_name
            if model_name not in model_times:
                model_times[model_name] = []
            model_times[model_name].append(result.processing_time)
        
        # 绘制箱线图
        models = list(model_times.keys())
        times_data = [model_times[model] for model in models]
        
        if times_data:
            bp = ax.boxplot(times_data, labels=models, patch_artist=True)
            
            # 设置SCI风格颜色
            for i, patch in enumerate(bp['boxes']):
                color = sci_colors[i % len(sci_colors)]
                patch.set_facecolor(color)
                patch.set_alpha(0.8)
                patch.set_edgecolor('black')
                patch.set_linewidth(0.8)
            
            # 设置其他箱线图元素的样式
            for element in ['whiskers', 'fliers', 'medians', 'caps']:
                plt.setp(bp[element], color='black', linewidth=1)
        
        ax.set_xlabel('Model', fontsize=11)
        ax.set_ylabel('Processing Time (seconds)', fontsize=11)
        ax.set_title('Rerank Only Time Distribution', fontsize=12, fontweight='bold', pad=15)
        ax.tick_params(axis='x', rotation=45, labelsize=10)
        ax.tick_params(axis='y', labelsize=10)
        ax.grid(True, alpha=0.3)
    
    def _plot_rerank_only_score_distribution(self, ax):
        """绘制rerank only分数分布图"""
        # SCI color scheme
        sci_colors = ['#2E86AB', '#A23B72', '#F18F01', '#C73E1D']
        
        if not self.rerank_only_results:
            ax.text(0.5, 0.5, 'No Data Available', ha='center', va='center', 
                   transform=ax.transAxes, fontsize=11)
            return
        
        # 按模型分组收集分数数据
        model_scores = {}
        for result in self.rerank_only_results:
            model_name = result.model_name
            if model_name not in model_scores:
                model_scores[model_name] = []
            
            # 从rerank_result中提取所有分数
            if hasattr(result, 'rerank_result') and result.rerank_result:
                scores = [score for _, score in result.rerank_result]
                model_scores[model_name].extend(scores)
            elif hasattr(result, 'top_scores') and result.top_scores:
                model_scores[model_name].extend(result.top_scores)
        
        # 绘制分数分布直方图
        for i, (model_name, scores) in enumerate(model_scores.items()):
            if scores and i < len(sci_colors):
                ax.hist(scores, alpha=0.7, label=model_name, bins=20, density=True,
                       color=sci_colors[i], edgecolor='black', linewidth=0.5)
        
        ax.set_title('Rerank Only Score Distribution', fontsize=12, fontweight='bold', pad=15)
        ax.set_xlabel('Score', fontsize=11)
        ax.set_ylabel('Density', fontsize=11)
        ax.legend(fontsize=9, loc='upper right')
        ax.tick_params(axis='both', which='major', labelsize=10)
        ax.grid(True, alpha=0.3)
    
    def _plot_rerank_only_performance(self, ax):
        """绘制rerank only模型性能对比图"""
        # SCI color scheme
        sci_colors = ['#2E86AB', '#A23B72', '#F18F01', '#C73E1D']
        
        if not self.rerank_only_results:
            ax.text(0.5, 0.5, 'No Data Available', ha='center', va='center', 
                   transform=ax.transAxes, fontsize=11)
            return
        
        # 按模型分组统计性能指标
        model_stats = {}
        for result in self.rerank_only_results:
            model_name = result.model_name
            if model_name not in model_stats:
                model_stats[model_name] = {
                    'times': [],
                    'scores': []
                }
            
            model_stats[model_name]['times'].append(result.processing_time)
            
            # 从rerank_result中提取所有分数
            if hasattr(result, 'rerank_result') and result.rerank_result:
                scores = [score for _, score in result.rerank_result]
                model_stats[model_name]['scores'].extend(scores)
            elif hasattr(result, 'top_scores') and result.top_scores:
                model_stats[model_name]['scores'].extend(result.top_scores)
        
        # 计算平均值
        models = list(model_stats.keys())
        avg_times = [np.mean(model_stats[model]['times']) for model in models]
        avg_scores = [np.mean(model_stats[model]['scores']) if model_stats[model]['scores'] else 0 for model in models]
        
        # 绘制柱状图
        x = np.arange(len(models))
        width = 0.35
        
        bars1 = ax.bar(x - width/2, avg_times, width, label='Avg Time (s)', 
                      color=sci_colors[0], alpha=0.8, edgecolor='black', linewidth=0.5)
        
        ax2 = ax.twinx()
        bars2 = ax2.bar(x + width/2, avg_scores, width, label='Avg Score', 
                       color=sci_colors[1], alpha=0.8, edgecolor='black', linewidth=0.5)
        
        # 添加数值标签
        for i, (time_val, score_val) in enumerate(zip(avg_times, avg_scores)):
            ax.text(i - width/2, time_val + 0.001, f'{time_val:.3f}', ha='center', va='bottom', 
                   fontsize=9, fontweight='bold')
            ax2.text(i + width/2, score_val + 0.01, f'{score_val:.3f}', ha='center', va='bottom', 
                    fontsize=9, fontweight='bold')
        
        ax.set_xlabel('Model', fontsize=11)
        ax.set_ylabel('Processing Time (seconds)', fontsize=11, color=sci_colors[0])
        ax2.set_ylabel('Average Score', fontsize=11, color=sci_colors[1])
        ax.set_title('Rerank Only Model Performance Comparison', fontsize=12, fontweight='bold', pad=15)
        ax.set_xticks(x)
        ax.set_xticklabels(models, rotation=45, ha='right', fontsize=10)
        
        # 设置图例
        lines1, labels1 = ax.get_legend_handles_labels()
        lines2, labels2 = ax2.get_legend_handles_labels()
        ax.legend(lines1 + lines2, labels1 + labels2, loc='upper right', fontsize=9)
        
        ax.tick_params(axis='both', which='major', labelsize=10)
        ax.grid(True, alpha=0.3)
    

    def _plot_rerank_only_consistency(self, ax):
        """绘制rerank only分数一致性分析图"""
        # SCI color scheme
        sci_colors = ['#2E86AB', '#A23B72', '#F18F01', '#C73E1D']
        
        if not self.rerank_only_results:
            ax.text(0.5, 0.5, 'No Data Available', ha='center', va='center', 
                   transform=ax.transAxes, fontsize=11)
            return
        
        # 按模型分组统计分数一致性
        model_consistency = {}
        for result in self.rerank_only_results:
            model_name = result.model_name
            if model_name not in model_consistency:
                model_consistency[model_name] = []
            
            if hasattr(result, 'score_variance') and result.score_variance is not None:
                # 使用已有的分数方差
                model_consistency[model_name].append(result.score_variance)
            elif hasattr(result, 'top_scores') and result.top_scores:
                # 计算top_scores的标准差作为一致性指标
                score_std = np.std(result.top_scores)
                model_consistency[model_name].append(score_std)
        
        # 计算平均一致性
        models = list(model_consistency.keys())
        avg_consistency = [np.mean(model_consistency[model]) if model_consistency[model] else 0 for model in models]
        
        if not models:
            ax.text(0.5, 0.5, 'No Consistency Data', ha='center', va='center', 
                   transform=ax.transAxes, fontsize=11)
            return
        
        # 绘制柱状图
        x = np.arange(len(models))
        bars = ax.bar(x, avg_consistency, color=[sci_colors[i % len(sci_colors)] for i in range(len(models))], 
                     alpha=0.8, edgecolor='black', linewidth=0.5)
        
        # 添加数值标签
        for i, val in enumerate(avg_consistency):
            ax.text(i, val + max(avg_consistency)*0.01 if avg_consistency else 0.001, 
                   f'{val:.4f}', ha='center', va='bottom', 
                   fontsize=9, fontweight='bold')
        
        ax.set_xlabel('Model', fontsize=11)
        ax.set_ylabel('Score Variance/Std Dev', fontsize=11)
        ax.set_title('Rerank Only Score Consistency Analysis\n(Lower is Better)', fontsize=12, fontweight='bold', pad=15)
        ax.set_xticks(x)
        ax.set_xticklabels(models, rotation=45, ha='right', fontsize=10)
        ax.tick_params(axis='both', which='major', labelsize=10)
        ax.grid(True, alpha=0.3)